"""
代码执行
"""

from mcp.server import FastMCP
import docker
from typing import Optional, List, Any, Dict
from pydantic import Field
import requests
import json
from .container import get_container
from docker.models.containers import Container
from pyflakes.api import check
from pyflakes.reporter import Reporter
import io


def get_serve_port(container: Container):
    for port, hosts in container.ports.items():
        if isinstance(hosts, list):
            for host in hosts:
                return host["HostPort"]
    raise Exception("该容器未开放服务端口.")


def add_code_tools(mcp: FastMCP, _docker: docker.DockerClient):

    @mcp.tool()
    async def install_package(
            name: str = Field(default="python-sandbox", description="容器名称."),
            container_id: Optional[str] = Field(default=None, description="容器id."),
            package_name: str = Field(default=..., description="python包名."),
    ) -> Dict[str, Any]:
        """为容器安装python包. """

        container = get_container(_docker, name, container_id)
        if container is None:
            result = {
                "success": False,
                "info": '',
                "error": '未找到相应的容器.'
            }
            return result

        port = get_serve_port(container)
        try:
            url = f'http://localhost:{port}/install'
            req = requests.post(url, data=json.dumps({"package_name": package_name}))
            return req.json()
        except Exception as e:
            result = {
                "success": False,
                "info": '',
                "error": str(e)
            }
        return result

    @mcp.tool()
    async def check_python_code(
            code: str = Field(default=..., description="python代码.")
    ):
        """静态检查 python代码
            Returns:
                以Json格式返回代码的错误个数error_count、警告信息warning_string、和错误信息error_string.
        """
        warning_stream = io.StringIO()
        error_stream = io.StringIO()
        reporter = Reporter(warningStream=warning_stream, errorStream=error_stream)
        error_count = check(code, filename="code", reporter=reporter)
        warning_string = warning_stream.getvalue()
        error_string = error_stream.getvalue()
        result = {
            "error_count": error_count,
            "warning_string": warning_string,
            "error_string": error_string
        }
        return result

    @mcp.tool()
    async def execute_python_code(
            name: str = Field(default="python-sandbox", description="python沙箱名称."),
            container_id: Optional[str] = Field(default=None, description="python沙箱id."),
            code: str = Field(default=..., description="python代码."),
            args: List = Field(default=[], description="python代码执行的列表参数."),
            kwargs: Dict = Field(default={}, description="python代码执行的字典参数."),
    ) -> Dict[str, Any]:
        """执行python代码"""

        container = get_container(_docker, name, container_id)
        if container is None:
            result = {
                'success': False,
                'data': '',
                'error': '未找到相应的容器.'
            }
            return result

        port = get_serve_port(container)
        try:
            url = f'http://localhost:{port}/run'
            req = requests.post(url, data=json.dumps({'code': code, 'args': args, 'kwargs': kwargs}))
            result = req.json()
        except Exception as e:
            result = {
                'success': False,
                'data': '',
                'error': str(e)
            }
        return result
